#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
弱口令漏洞扫描脚本
检测目标服务是否存在弱口令问题
"""

import requests
import time
import re
import socket
from typing import Dict, Any, List
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 尝试导入paramiko，如果失败则设置为None
paramiko = None
try:
    import paramiko
    logger.info("成功导入paramiko模块")
except ImportError as e:
    logger.warning(f"无法导入paramiko模块: {str(e)}")
    logger.warning("SSH服务弱口令扫描功能将不可用")

# 脚本元数据
NAME = "弱口令漏洞扫描器"
DESCRIPTION = "检测常见服务是否存在弱口令问题，包括HTTP表单、SSH、FTP等服务"
AUTHOR = "Security Team"
VERSION = "1.0.0"
VULNERABILITY_TYPE = "weak_password"
SEVERITY = "high"
ENABLED = True

# 必要参数
REQUIRED_PARAMS = ["target"]

# 可选参数及默认值
OPTIONAL_PARAMS = {
    "scan_depth": 2,
    "timeout": 10,
    "max_retries": 3,
    "port": 80,
    "service_type": "http"
}

# 弱口令字典（简单版本，实际应用中可以加载更大的字典文件）
DEFAULT_CREDENTIALS = [
    ("admin", "admin"),
    ("admin", "password"),
    ("admin", "123456"),
    ("admin", "12345678"),
    ("admin", "qwerty"),
    ("root", "root"),
    ("root", "toor"),
    ("root", "123456"),
    ("root", "12345678"),
    ("user", "user"),
    ("test", "test"),
    ("guest", "guest"),
    ("default", "default")
]

# 常见登录表单信息
LOGIN_FORM_PATTERNS = [
    {
        "name": "Generic Form 1",
        "username_field": "username",
        "password_field": "password",
        "login_paths": ["/login", "/signin", "/admin", "/wp-login.php", "/login.php"],
        "success_pattern": r"<title>.*Dashboard.*</title>|Welcome.*|Dashboard|admin|dashboard",
        "failure_pattern": r"Invalid username|Invalid password|登录失败|用户名或密码错误"
    },
    {
        "name": "Generic Form 2",
        "username_field": "user",
        "password_field": "pass",
        "login_paths": ["/login", "/signin", "/admin", "/wp-login.php", "/login.php"],
        "success_pattern": r"<title>.*Dashboard.*</title>|Welcome.*|Dashboard|admin|dashboard",
        "failure_pattern": r"Invalid username|Invalid password|登录失败|用户名或密码错误"
    }
]

def check_ssh(target: str, port: int, username: str, password: str, timeout: int) -> bool:
    """
    检查SSH服务弱口令
    """
    # 检查paramiko是否可用
    if paramiko is None:
        logger.warning("paramiko模块不可用，无法执行SSH弱口令检查")
        return False
    
    try:
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(target, port=port, username=username, password=password, timeout=timeout)
        client.close()
        return True
    except paramiko.AuthenticationException:
        return False
    except (paramiko.SSHException, socket.error, Exception):
        return False

def check_ftp(target: str, port: int, username: str, password: str, timeout: int) -> bool:
    """
    检查FTP服务弱口令
    """
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(timeout)
        sock.connect((target, port))
        
        # 接收欢迎信息
        sock.recv(1024)
        
        # 发送用户名
        sock.sendall(f"USER {username}\r\n".encode())
        sock.recv(1024)
        
        # 发送密码
        sock.sendall(f"PASS {password}\r\n".encode())
        response = sock.recv(1024).decode()
        
        sock.close()
        
        # 成功登录通常返回230
        return response.startswith("230")
    except Exception:
        return False

def check_http_login(target: str, login_path: str, form_config: dict, username: str, password: str, timeout: int) -> Dict[str, Any]:
    """
    检查HTTP表单登录弱口令
    """
    try:
        login_url = f"{target}{login_path}" if target.endswith("/") else f"{target}{login_path}"
        
        # 先获取登录页面（可能需要获取CSRF token等）
        session = requests.Session()
        response = session.get(login_url, timeout=timeout)
        
        # 准备登录数据
        data = {
            form_config["username_field"]: username,
            form_config["password_field"]: password
        }
        
        # 尝试提取CSRF token
        csrf_patterns = [
            r'<input[^>]*name=["\']csrf_token["\'][^>]*value=["\']([^"\']+)["\']',
            r'<input[^>]*name=["\']token["\'][^>]*value=["\']([^"\']+)["\']',
            r'<input[^>]*name=["\']_token["\'][^>]*value=["\']([^"\']+)["\']'
        ]
        
        for pattern in csrf_patterns:
            match = re.search(pattern, response.text)
            if match:
                token_name = re.search(r'name=["\']([^"\']+)["\']', match.group(0)).group(1)
                data[token_name] = match.group(1)
                break
        
        # 发送登录请求
        login_response = session.post(login_url, data=data, timeout=timeout, allow_redirects=True)
        
        # 检查是否登录成功
        success = False
        failure = False
        
        if re.search(form_config["success_pattern"], login_response.text, re.IGNORECASE):
            success = True
        elif re.search(form_config["failure_pattern"], login_response.text, re.IGNORECASE):
            failure = True
        elif "logout" in login_response.text.lower():
            success = True
        elif login_response.status_code == 302:
            # 检查重定向后的页面
            redirect_response = session.get(login_response.url, timeout=timeout)
            if re.search(form_config["success_pattern"], redirect_response.text, re.IGNORECASE):
                success = True
        
        return {
            "success": success,
            "failure": failure,
            "status_code": login_response.status_code
        }
    except Exception as e:
        logger.error(f"HTTP登录测试失败: {str(e)}")
        return {
            "success": False,
            "failure": True,
            "status_code": 0,
            "error": str(e)
        }

def scan(target: str, **kwargs) -> Dict[str, Any]:
    """
    执行弱口令漏洞扫描
    
    Args:
        target: 目标URL或IP
        **kwargs: 可选参数
            scan_depth: 扫描深度 (1-5)
            timeout: 请求超时时间
            max_retries: 最大重试次数
            port: 目标端口
            service_type: 服务类型 (http, ssh, ftp, all)
    """
    start_time = time.time()
    results = {
        "target": target,
        "vulnerability_type": VULNERABILITY_TYPE,
        "severity": SEVERITY,
        "scan_depth": kwargs.get("scan_depth", 2),
        "start_time": start_time,
        "end_time": None,
        "duration": None,
        "vulnerable": False,
        "vulnerabilities": [],
        "details": {
            "total_tests": 0,
            "passed_tests": 0,
            "failed_tests": 0,
            "test_details": []
        }
    }
    
    # 获取参数
    timeout = kwargs.get("timeout", 10)
    max_retries = kwargs.get("max_retries", 3)
    scan_depth = kwargs.get("scan_depth", 2)
    port = kwargs.get("port", 80)
    service_type = kwargs.get("service_type", "http")
    
    # 根据扫描深度调整密码字典大小
    if scan_depth < 3:
        test_credentials = DEFAULT_CREDENTIALS[:6]  # 只测试最常见的6个
    else:
        test_credentials = DEFAULT_CREDENTIALS
    
    logger.info(f"开始弱口令扫描: {target} (类型: {service_type}, 端口: {port})")
    
    # 要测试的服务
    services_to_test = []
    if service_type == "all" or service_type == "http":
        services_to_test.append("http")
    if service_type == "all" or service_type == "ssh":
        services_to_test.append("ssh")
    if service_type == "all" or service_type == "ftp":
        services_to_test.append("ftp")
    
    # 测试每个服务
    for service in services_to_test:
        service_results = {
            "service": service,
            "port": port if service == "http" else {
                "ssh": 22,
                "ftp": 21
            }.get(service, port),
            "vulnerable": False,
            "found_credentials": [],
            "tested_credentials": 0
        }
        
        # 测试弱口令
        for username, password in test_credentials:
            results["details"]["total_tests"] += 1
            service_results["tested_credentials"] += 1
            
            success = False
            test_result = {
                "service": service,
                "username": username,
                "password": password,
                "success": False
            }
            
            try:
                if service == "ssh":
                    success = check_ssh(target, 22, username, password, timeout)
                elif service == "ftp":
                    success = check_ftp(target, 21, username, password, timeout)
                elif service == "http":
                    # 尝试不同的登录路径和表单配置
                    for login_path in ["/login", "/signin", "/admin", "/wp-login.php", "/login.php"]:
                        for form_config in LOGIN_FORM_PATTERNS[:1]:  # 先只测试一个表单配置
                            http_result = check_http_login(target, login_path, form_config, username, password, timeout)
                            if http_result["success"]:
                                success = True
                                test_result["login_path"] = login_path
                                break
                        if success:
                            break
                
                if success:
                    service_results["vulnerable"] = True
                    service_results["found_credentials"].append((username, password))
                    test_result["success"] = True
                    
                    # 添加到漏洞列表
                    vulnerability = {
                        "service": service,
                        "port": service_results["port"],
                        "username": username,
                        "password": password,
                        "recommendation": f"为{service}服务设置强密码，启用双因素认证，限制登录尝试次数"
                    }
                    if service == "http" and "login_path" in test_result:
                        vulnerability["login_path"] = test_result["login_path"]
                    
                    results["vulnerabilities"].append(vulnerability)
                    results["vulnerable"] = True
                
                results["details"]["passed_tests"] += 1
            except Exception as e:
                logger.error(f"测试{service}服务弱口令失败: {str(e)}")
                results["details"]["failed_tests"] += 1
                test_result["error"] = str(e)
            
            results["details"]["test_details"].append(test_result)
        
        # 避免过度测试，如果发现了弱口令，可以提前结束
        if service_results["vulnerable"] and scan_depth < 4:
            logger.info(f"在{service}服务上发现弱口令，暂停该服务的进一步测试")
            break
    
    # 计算结束时间和持续时间
    end_time = time.time()
    results["end_time"] = end_time
    results["duration"] = end_time - start_time
    
    # 生成总结报告
    if results["vulnerable"]:
        logger.warning(f"在 {target} 发现 {len(results['vulnerabilities'])} 个弱口令漏洞")
        results["summary"] = f"发现{len(results['vulnerabilities'])}个弱口令漏洞，严重威胁系统安全"
    else:
        logger.info(f"{target} 未发现明显的弱口令漏洞")
        results["summary"] = "未发现明显的弱口令漏洞，但仍建议定期更新密码并使用强密码策略"
    
    return results